package tools

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"encoding/base64"
	"errors"
	"fmt"
	"math/rand"
)

var AesKey = "52FDFC072182654F163F5F0F9A621D729566C74D10037C4D7BBB0407D1E2C649" // 32位, 如有需要使用 RSA加密AES密钥

func encryptAES(message string, key []byte, ivOpt ...[]byte) (ciphertextBase64 string, err error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", fmt.Errorf("error creating new cipher: %w", err)
	}

	var iv []byte
	switch len(ivOpt) {
	case 0:
		// Generate a fresh Initialization Vector each time when none was passed explicitly by caller.
		iv, err = generateRandomBytes(aes.BlockSize)
		if err != nil {
			return "", fmt.Errorf("could not create initialization vector: %v", err.Error())
		}
	default:
		iv = ivOpt[0]

		// Ensure that user-provided IV has correct length matching chosen block size of selected AES variant.
		expectedLen := aes.BlockSize
		actualLen := len(iv)
		if expectedLen != actualLen {
			errMsg := fmt.Sprintf(
				"invalid initialization vector length (%d), must match blocksize=%d",
				actualLen, expectedLen,
			)
			return "", errors.New(errMsg)
		}
	}

	cfbMode := cipher.NewCBCEncrypter(block, iv)
	paddingLength := aes.BlockSize - len(message)%aes.BlockSize
	paddedMessage := append([]byte(message),
		bytes.Repeat([]byte{byte(paddingLength)}, paddingLength)...,
	)

	buffer := make([]byte, len(paddedMessage)+len(iv)) // Reserve space also for storing IV alongside ciphertext.

	copy(buffer[:aes.BlockSize], iv) // Store IV first...
	cfbMode.CryptBlocks(buffer[aes.BlockSize:], paddedMessage)

	return base64.StdEncoding.EncodeToString(buffer), nil
}
func decryptAES(ciphertextBase64 string, key []byte) (plainText string, err error) {
	data, err := base64.StdEncoding.DecodeString(ciphertextBase64)
	if err != nil {
		return "", fmt.Errorf("decoding failed due to malformed input: %w", err)
	}

	if len(data) < aes.BlockSize { // Check whether there's enough room left after extracting IV part out.
		return "", errors.New("malformed payload detected; too short to contain valid initialization vector segment")
	}

	block, err := aes.NewCipher(key)
	if err != nil {
		return "", fmt.Errorf("cipher creation failure encountered while attempting decryption step: %w", err)
	}

	iv := data[:aes.BlockSize]      // Extract stored IV value which precedes our target ciphertext proper.
	payload := data[aes.BlockSize:] //

	mode := cipher.NewCBCDecrypter(block, iv)
	mode.CryptBlocks(payload, payload)

	unPaddingValue := int(payload[len(payload)-1])
	if unPaddingValue > aes.BlockSize || unPaddingValue == 0 {
		return "", errors.New("incorrectly formatted padding found upon removing trailing garbage characters appended earlier during encryption phase")
	} else if !bytes.Equal(bytes.Repeat([]byte{byte(unPaddingValue)},
		unPaddingValue), payload[len(payload)-unPaddingValue:]) {
		return "", errors.New("inconsistent/unexpected padding sequence discovered near end-of-message boundary area post-decrypting attempt")
	}

	result := payload[:len(payload)-unPaddingValue]
	return string(result), nil
}

// 以此方法可以生成随机字节
func generateRandomBytes(keySize int) ([]byte, error) {
	key := make([]byte, keySize)
	_, err := rand.Read(key)
	if err != nil {
		return nil, err
	}
	return key, nil
}

func EncryptAES(baseContent string) (res string, err error) {
	key, _ := HexStrToByteArray(AesKey)
	res, err = encryptAES(baseContent, key)
	if err != nil {
		return
	}
	return
}
func DecryptAES(baseContent string) (res string, err error) {
	key, _ := HexStrToByteArray(AesKey)
	res, err = decryptAES(baseContent, key)
	if err != nil {
		return
	}
	return
}
